#include "SymmetryDetector.h"

#include "MeanShift.h"
#include "MeanShift.cpp"

#include "ICP.h"


#define	BASE_NEIGHBOR_RADIUS	0.05
#define BASE_NUM_CLOSEST		20
#define DEGREE_POLYNOMIAL		2
#define DEGREE_MONGE			2


SymmetryDetector::SymmetryDetector( SurfaceMeshModel* src_mesh)
{
	this->mesh = src_mesh;

	// set parameters
	numFeaturePoints = 500;
	numBasePoints = 10000;
	pairingThreshold = 1.0;
	distinctiveRatio = 0.05;

	// show options
	showPoints = false;
	showPointSignatures = false;
	showPrunedSignatures = false;

	// Debug
	debug_points.clear();
	debug_points.push_back(PointSoup(Qt::cyan));
	debug_points.push_back(PointSoup(Qt::red));
	debug_points.push_back(PointSoup(Qt::green));
	debug_points.push_back(PointSoup(Qt::blue));

	debug_vectors = VectorSoup(Qt::magenta);
}

SymmetryDetector::~SymmetryDetector()
{
	delete feature_kdtree;
	delete base_kdtree;
}

void SymmetryDetector::draw()
{
	if (showPoints)				drawPoints();
	if (showPointSignatures)	drawPointSignatures();
	if (showPrunedSignatures)	drawPrunedSignatures();

	// Debug
	drawDebug();
}

void SymmetryDetector::drawPoints()
{
	PointSoup feature_ps(Qt::yellow);
	foreach (Sample s, feature_samples) feature_ps.addPoint(s.p);
	feature_ps.draw();

	PointSoup base_ps(Qt::cyan, 4.0);
	foreach (Sample s, base_samples) base_ps.addPoint(s.p);
	base_ps.draw();
}

void SymmetryDetector::drawPointSignatures()
{
	double max_curvature = -1;
	foreach (PointSignature sig, point_signatures)
	{
		double curr = abs(sig.k1);
		if ( curr > max_curvature )	max_curvature = curr;
	}

	VectorSoup normals(Qt::green);
	LineSegments maxCurv(Qt::red), minCurv(Qt::blue);

	double max_len = 0.02, len1, len2;
	Vec3d sp;
	std::vector<Vec3d> na, nb, d1a, d1b, d2a, d2b;
	foreach (PointSignature sig, point_signatures)
	{
		//if (sig.type == UMBILICAL) continue;

		normals.addVector(sig.p, max_len * sig.n);

		// Shift up along the normal for better visualization
		sp = sig.p + max_len / 10 * sig.n;

		len1 = sqrt(abs(sig.k1 / max_curvature)) * max_len;
		maxCurv.addLine(sp - len1 * sig.d1, sp + len1 * sig.d1);

		len2 = sqrt(abs(sig.k2 / max_curvature)) * max_len;
		minCurv.addLine(sp - len2 * sig.d2, sp + len2 * sig.d2);
	}

	//normals.draw();
	maxCurv.draw();
	minCurv.draw();
}

void SymmetryDetector::drawPrunedSignatures()
{
	double max_curvature = -1;
	foreach (PointSignature sig, pruned_point_signatures)
	{
		double curr = abs(sig.k1);
		if ( curr > max_curvature )	max_curvature = curr;
	}

	VectorSoup normals(Qt::green);
	LineSegments maxCurv(Qt::red), minCurv(Qt::blue);

	double max_len = 0.02, len1, len2;
	Vec3d sp;
	std::vector<Vec3d> na, nb, d1a, d1b, d2a, d2b;
	foreach (PointSignature sig, pruned_point_signatures)
	{
		//if (sig.type == UMBILICAL) continue;

		normals.addVector(sig.p, max_len * sig.n);

		// Shift up along the normal for better visualization
		sp = sig.p + max_len / 10 * sig.n;

		len1 = sqrt(abs(sig.k1 / max_curvature)) * max_len;
		maxCurv.addLine(sp - len1 * sig.d1, sp + len1 * sig.d1);

		len2 = sqrt(abs(sig.k2 / max_curvature)) * max_len;
		minCurv.addLine(sp - len2 * sig.d2, sp + len2 * sig.d2);
	}

	//normals.draw();
	maxCurv.draw();
	minCurv.draw();
}


void SymmetryDetector::showMessage( const char* format, ... )
{
	const size_t buffer_length=256;
	char buffer[buffer_length];
	va_list args;
	va_start (args, format);
	vsnprintf(buffer,buffer_length,format, args);
	va_end (args);

	emit(showStatusMessage(buffer));

}


void SymmetryDetector::sampleGeometry()
{
	// Sample Geometry
	Sampler sampler(mesh);
	feature_samples = sampler.getSamples(numFeaturePoints);
	base_samples = sampler.getSamples(numBasePoints);

	// Construct the kdtrees
	feature_kdtree = new NanoflannKdTree3d;
	foreach (Sample s, feature_samples) feature_kdtree->addPoint(s.p);
	feature_kdtree->build();

	base_kdtree = new NanoflannKdTree3d;
	foreach (Sample s, base_samples) base_kdtree->addPoint(s.p);
	base_kdtree->build();
}

std::vector<Vec3d> SymmetryDetector::closestBasePoints( Vec3d query_point )
{
	std::vector<Vec3d> closest_base_points;

	KDResults base_kd_results;
	base_kdtree->ball_search(query_point, BASE_NEIGHBOR_RADIUS, base_kd_results);
	
	// force to collect NUM_CLOSEST points
	if (base_kd_results.size() < BASE_NUM_CLOSEST)
	{
		base_kd_results.clear();
		base_kdtree->k_closest(query_point, BASE_NUM_CLOSEST, base_kd_results);
	}

	for (int i = 0; i < (int) base_kd_results.size(); i++)
	{
		size_t base_id = base_kd_results[i].first;
		closest_base_points.push_back(base_samples[base_id].p);
	}

	return closest_base_points;
}

Monge_form SymmetryDetector::mongeFormAtFeaturePoint( int fid )
{
	// Collect neighbors
	Vec3d query_point = feature_samples[fid].p;
	std::vector<Vec3d> closeat_base_points = closestBasePoints(query_point);

	// The Monge form will be estimated on the first point
	std::vector<Vec3d> closest_points;
	closest_points.push_back(query_point);
	closest_points.insert(closest_points.end(), closeat_base_points.begin(), closeat_base_points.end());
		
	// Monge fitting
	Monge_form monge_form;
	Monge_via_jet_fitting monge_fit;
	monge_form = monge_fit(closest_points.begin(), closest_points.end(), DEGREE_POLYNOMIAL, DEGREE_MONGE);

	return monge_form;
}

void SymmetryDetector::computePointSignatures()
{
	point_signatures.clear();

	for (int i = 0; i < (int) feature_samples.size(); i++)
	{
		Monge_form monge_form = mongeFormAtFeaturePoint(i);

		// Use the input normals to adjust Monge forms
		monge_form.comply_wrt_given_normal(feature_samples[i].n);

		// point signatures
		point_signatures.push_back(PointSignature(i, monge_form));
	}
}

void SymmetryDetector::pruneFeaturePoints()
{
	pruned_point_signatures.clear();

    foreach (PointSignature sig, point_signatures)
    {
		sig.determinType();

		if (sig.type != UMBILICAL)
			pruned_point_signatures.push_back(sig);
    }
}

void SymmetryDetector::constructPairingKdtree()
{
	// Construct the kdtree for pairing
	pairing_kdtree = new NanoflannKdTree2d;
	foreach (PointSignature ps, pruned_point_signatures)
	{
		pairing_kdtree->addPoint(ps.curvaturesAsVec2d());
	}
	pairing_kdtree->build();
}

std::vector<int> SymmetryDetector::peerPFIndices(int seed_id)
{
	PointSignature seed = pruned_point_signatures[seed_id];
	Vec2d query_point = seed.curvaturesAsVec2d();

	KDResults kd_results;
	pairing_kdtree->ball_search(query_point, pairingThreshold, kd_results);

	// Skip the first one which is the seed itself
	std::vector<int> peers;
	for (int i = 1; i < (int) kd_results.size(); i++)
		peers.push_back(kd_results[i].first);

	return peers;
}

void SymmetryDetector::generateVoters()
{
	sym_voters.clear();

	// Randomly select a subset as seeds
	int numPoints =  pruned_point_signatures.size();
	int numSeeds = numPoints / 5;
	QSet<int> seed_ids = subset_indices(numSeeds, numPoints);

	// Pair each seed with its peers
	int distictive_num = distinctiveRatio * pruned_point_signatures.size();
	foreach (int seed_id, seed_ids)
	{
		std::vector<int> peer_ids = peerPFIndices(seed_id);

		// The feature has to be distinctive
		// General features vote for everything and thus very costly
		if (peer_ids.size() > distictive_num) continue;

		foreach(int peer_id, peer_ids)
		{
			PointSignature seed = pruned_point_signatures[seed_id];
			PointSignature peer = pruned_point_signatures[peer_id];
			sym_voters.push_back(SymmetryVoter(seed, peer));
		}
	}
}

void SymmetryDetector::prepareForDetection()
{
	showMessage("Sampling Geometry...");
	this->sampleGeometry();

	showMessage("Computing point signatures...");
	this->computePointSignatures();

	showMessage("Pruning point signatures...");
	this->pruneFeaturePoints();
	showMessage("Pruning point signatures... Done: # sig = %d", pruned_point_signatures.size());

	showMessage("Constructing pairing kdtree...");
	this->constructPairingKdtree();

	showMessage("Generating symmetry voters...");
	this->generateVoters();
	showMessage("Generating symmetry voters... Done: # voters = %d", sym_voters.size());
}


void SymmetryDetector::clusterReflections()
{
	// points in transformation space
	std::vector<Vec4d> reflections;
	Vec4d ref;
	foreach (SymmetryVoter voter, sym_voters)
	{
		if (voter.reflection(ref))
		{
			// there are two formulations for the same plane
			reflections.push_back( ref);
			reflections.push_back(-ref);
		}
	}

	// detect modes in transformation space
	showMessage("Clustering via meanshift...");
	MeanShift<4> ref_meanshift(reflections);
	ref_meanshift.cluster();

	// modes come in pairs
	std::vector<Vec4d> &ref_modes = ref_meanshift.modes;
	double diff_threshold = 0.05;
	std::vector<bool> is_redundant(ref_modes.size(), false);
	for (int i = 0; i < (int) ref_modes.size(); i++)
	{
		if (is_redundant[i]) continue;

		// a new mode and its voters
		ref_planes.push_back(ref_modes[i]);
		ref_density.push_back(ref_meanshift.densityAt(ref_planes.back()));
		QSet<int> pids = ref_meanshift.pointIDsNearMode(i);

		// eliminate its peer if exists
		for (int j = i+1; j < (int)ref_modes.size(); j++)
		{
			double diff = (ref_modes[i] + ref_modes[j]).norm();

			if (diff < diff_threshold)
			{
				is_redundant[j] = true;
				ref_density.back() += ref_meanshift.densityAt(ref_modes[j]);
				pids += ref_meanshift.pointIDsNearMode(j);
			}
		}

		// Each voter generates two points in transformation space
		QSet<int> voterIDs;
		foreach (int pid, pids) voterIDs.insert( pid / 2 );

		ref_voterIDs.push_back(voterIDs);
	}
}

void SymmetryDetector::verifyReflection(int i)
{
	// copy the original mesh
	QString new_name = mesh->name + "_reflected";
	SurfaceMeshModel *reflected_mesh = new SurfaceMeshModel(new_name, new_name);
	reflected_mesh->deep_copy(*mesh);
	reflected_mesh->color = QColor(186, 172, 96);
	emit(addModelToDocument(reflected_mesh));
	mesh->color = QColor(86, 117, 147);

	// reflect the copied mesh
	Vec4d ref_plane(1,0,0,0);// = ref_planes[i];
	ref_planes.push_back(ref_plane); ref_density.push_back(1);
	visualizeReflections();
	Vector3VertexProperty reflected_points = reflected_mesh->vertex_coordinates();
	Vector3VertexProperty reflected_normals = reflected_mesh->vertex_normals();
	foreach(Vertex v, reflected_mesh->vertices())
	{
		reflected_points[v] = reflect(reflected_points[v], ref_plane);
		reflected_normals[v] = reflect(reflected_normals[v], ref_plane);
	}

	// region growing 
	ICP icp(reflected_mesh, mesh);
	NanoflannKdTree3d *kdtree = icp.target_kdtree;		// reuse the kdtree for efficiency
	QSet<int> &voterIDs = ref_voterIDs[i];
	foreach (int vi, voterIDs)
	{
		// the voter on the positive side of reflection plane
		SymmetryVoter &voter = sym_voters[vi];
		int fp_id;
		if (onPositiveSide(voter.firstPoint(), ref_plane))
			fp_id = voter.firstID();
		else
			fp_id = voter.secondID();


		// grow the region
		QStack<Vertex> activeVertices;
		BoolVertexProperty isVisited = reflected_mesh->vertex_property("isVisited", false);

		if (feature_samples[fp_id].isVertex)
		{
			Vertex v = feature_samples[fp_id].vertex();
			activeVertices.push_back(v);
		}
		else
		{
			Face f = feature_samples[fp_id].face();
			Vertex_around_face_circulator fvit, fvend;
			fvit = fvend = vertices(f);
			do{
				activeVertices.push_back(Vertex(fvit));
			} while (++fvit != fvend);
		}

		foreach (Vertex v, activeVertices)
			isVisited[v] = true;

	}
}

void SymmetryDetector::detectReflections()
{
	// initial
	ref_planes.clear();
	ref_voterIDs.clear();

	// prepare voters
	this->prepareForDetection();
	if (sym_voters.empty()) return;

	// cluster in transformation space
	showMessage("Clustering via meanshift...");
	this->clusterReflections();
	showMessage("Clustering via meanshift... Done: # ref_planes = %d", ref_planes.size());

	// verify the resulted modes
	this->verifyReflection(0);
}



////////////////////////////// DEBUG ////////////////////////////////////////////

void SymmetryDetector::test()
{
	//verifyReflection(0);
	int n = mesh->n_vertices();
	mesh->remove_vertex(Vertex(0));
	n = mesh->n_vertices(); 
}

void SymmetryDetector::debug()
{
	visualizePeers();
}


void SymmetryDetector::drawDebug()
{
	foreach(PointSoup ps, debug_points) ps.draw();
	debug_vectors.draw();
	debug_planes.draw();
}

void SymmetryDetector::visualizePeers()
{
	int numPoints = pruned_point_signatures.size();
	int seed_id = uniform_int(0, numPoints - 1);
	PointSignature seed = pruned_point_signatures[seed_id];

	debug_points[3].clear();
	debug_points[3].addPoint(seed.p);

	std::vector<int> peer_ids = peerPFIndices(seed_id);

	debug_points[2].clear();
	foreach (int id, peer_ids)
	{
		PointSignature peer = pruned_point_signatures[id];
		debug_points[2].addPoint(peer.p);
	}

}

void SymmetryDetector::visualize_kd_result( Vec3d query_point )
{
	debug_points[1].addPoint(query_point);
	foreach(Vec3d p, closestBasePoints(query_point)) 
		debug_points[2].addPoint(p);
}



void SymmetryDetector::visualizePointTypes()
{
	foreach(PointSignature sig, point_signatures)
	{
		sig.determinType();
		debug_points[sig.type].addPoint(sig.p);

	}
}

void SymmetryDetector::addSignatureProperty()
{
	// Use info plugin to visualize

	using namespace SurfaceMesh;

	IntegerVertexProperty vType = mesh->vertex_property("v:type", 0);
	Vector3VertexProperty vCurvature = mesh->vertex_property("v:curvs", Vec3d(0.0));
	for (int i = 0; i < (int) point_signatures.size(); i++)
	{
		Vertex v(i);
		PointSignature sig = point_signatures[i];

		vType[v] = sig.type;
		vCurvature[v] = Vec3d(sig.k1, sig.k2, sig.K);
	}
}

void SymmetryDetector::testPeers()
{
	//visualize_kd_result(feature_samples[0].p);
	computePointSignatures();
	pruneFeaturePoints();
	constructPairingKdtree();
}

void SymmetryDetector::visualizeReflections()
{
	debug_planes.clear();

	double max_density = -1;
	foreach (double rho, ref_density)
	{
		if (max_density < rho)
			max_density = rho;
	}

	for (size_t i = 0; i < ref_planes.size(); i++)
	{
		Vec4d pm = ref_planes[i];
		Vec3d n(pm[0], pm[1], pm[2]);
		Vec3d p = -pm[3] * n;

		double s = ref_density[i] / max_density;
		debug_planes.addPlane(p, n, qtColdColor(s));
	}
}
